import numpy as np
import utils
from astropy.time import Time
import astropy.units as u

j_2 = 0.00108263
A_2 = 1.5 * j_2
a_e = 6378137
GE = 3986004418e5
time_unit = np.sqrt(a_e ** 3 / GE)

class MeanElements:
    def __init__(self, sigma=np.zeros(6)):
        self.sigma = sigma
        self.sigma[5] = self.sigma[5] % (2 * np.pi)
        self.a = sigma[0]
        self.e = sigma[1]
        self.i = sigma[2]
        self.omega = sigma[3]
        self.w = sigma[4]
        self.M = sigma[5]

        self.E = utils.newton_generation(self.M, self.e, self.M, 1e-12)
        self.r = self.a * (1 - self.e * np.cos(self.E))
        self.p = self.a * (1 - self.e ** 2)
        self.n = self.a ** (-3/2)
        # self.f = 2 * np.arctan(np.sqrt(1 + self.e) / (1 - self.e)) * np.tan(self.E / 2)
        self.f = np.arctan2(np.sqrt(1 - self.e ** 2) * np.sin(self.E), np.cos(self.E) - self.e) % (2 * np.pi)

    def get_1st_short_and_long(self):
        """
        由某时刻的根数计算一阶长短周期项
        :return: 一阶短周期项 + 一阶长周期项
        """
        # 计算1阶长周期项，参考书本（4.165-4.170）
        a_long_1st = 0
        i_long_1st = -1 / (2 * self.p * self.p) * np.sin(2 * self.i) / (4 - 5 * (np.sin(self.i) ** 2)) * A_2 * (7/12 - 5/8 * (np.sin(self.i) ** 2)) * (self.e ** 2) * np.cos(2 * self.w)
        omega_long_1st = (-1 / (self.p * self.p) * np.cos(self.i) / ((4 - 5 * (np.sin(self.i) ** 2)) ** 2) * A_2 * (7/3 - 5 * (np.sin(self.i) ** 2) + 25/8 * (np.sin(self.i) ** 4)) * (self.e ** 2) * np.sin(2 * self.w))
        e_long_1st = -(1 - self.e ** 2) / self.e * np.tan(self.i) * i_long_1st
        w_long_1st = -1 / (self.p * self.p) / ((4 - 5 * (np.sin(self.i) ** 2)) ** 2) * (A_2 * ((np.sin(self.i) ** 2) * (25/3 - 245/12 * (np.sin(self.i) ** 2) + 25/2 * (np.sin(self.i) ** 4)) - (self.e ** 2) * (7/3 - 17/2 * (np.sin(self.i) ** 2) + 65/6 * (np.sin(self.i) ** 4) - 75/16 * (np.sin(self.i) ** 6)))) * np.sin(2 * self.w)
        M_long_1st = 1 / (self.p * self.p) * np.sqrt(1 - self.e ** 2) / (4 - 5 * (np.sin(self.i) ** 2)) * (np.sin(self.i) ** 2) * (A_2 * ((25/12 - 5/2 * (np.sin(self.i) ** 2)) - (self.e ** 2) * (7/12 - 5/8 * (np.sin(self.i) ** 2)))) * np.sin(2 * self.w)

        sigma_long_1st = [a_long_1st, e_long_1st, i_long_1st, omega_long_1st, w_long_1st, M_long_1st]

        # 计算1阶短周期项，参考书本（4.159-4.164）
        a_short_1st = A_2 / self.a * (2/3 * (1 - 3/2 * (np.sin(self.i) ** 2)) * (((self.a / self.r) ** 3) - ((1 - self.e * self.e) ** (-3/2))) + (np.sin(self.i) ** 2) * ((self.a / self.r) ** 3) * np.cos(2 * (self.f + self.w)))
        i_short_1st = A_2 / (4 * self.p * self.p) * np.sin(2 * self.i) * (self.e * np.cos(self.f + 2 * self.w) + np.cos(2 * (self.f + self.w)) + self.e / 3 * np.cos(3 * self.f + 2 * self.w))
        e_short_1st = (1 - self.e * self.e) / self.e * (1 / (2 * self.a) * a_short_1st - (np.tan(self.i) * i_short_1st))
        omega_short_1st = -A_2 / (self.p * self.p) * np.cos(self.i) * ((self.f - self.M + self.e * np.sin(self.f)) - 1/2 * (self.e * np.sin(self.f + 2 * self.w) + np.sin(2 * (self.f + self.w)) + self.e / 3 * np.sin(3 * self.f + 2 * self.w)))
        w_short_1st = -np.cos(self.i) * omega_short_1st + A_2 / (self.p * self.p) * ((1 - 3/2 * (np.sin(self.i) ** 2)) * (self.f - self.M + self.e * np.sin(self.f) + (1 / self.e - self.e / 4) * np.sin(self.f) + 1/2 * np.sin(2 * self.f) + self.e / 12 * np.sin(3 * self.f)) + (np.sin(self.i) ** 2) * (-(1 / (4 * self.e) - 7/16 * self.e) * np.sin(self.f + 2 * self.w) + 3/4 * np.sin(2 * (self.f + self.w)) + (7 / (12 * self.e) + 11/48 * self.e) * np.sin(3 * self.f + 2 * self.w) + 3/8 * np.sin(4 * self.f + 2 * self.w) + self.e / 16 * np.sin(5 * self.f + 2 * self.w) + self.e / 16 * np.sin(self.f - 2 * self.w)))
        M_short_1st = A_2 / (self.p ** 2) * np.sqrt(1 - self.e ** 2) * (-(1 - 3/2 * (np.sin(self.i) ** 2)) * ((1 / self.e - self.e / 4) * np.sin(self.f) + 1/2 * np.sin(2 * self.f) + self.e / 12 * np.sin(3 * self.f)) - (np.sin(self.i) ** 2) * (-(1 / (4 * self.e) + 5/16 * self.e) * np.sin(self.f + 2 * self.w) + (7 / (12 * self.e) - 1/48 * self.e) * np.sin(3 * self.f + 2 * self.w) + 3/8 * np.sin(4 * self.f + 2 * self.w) + self.e / 16 * np.sin(5 * self.f + 2 * self.w) + self.e / 16 * np.sin(self.f - 2 * self.w)))

        sigma_short_1st = np.array([a_short_1st, e_short_1st, i_short_1st, omega_short_1st, w_short_1st, M_short_1st])

        # mean_sigma = self.sigma - sigma_short_1st - sigma_long_1st
        return sigma_short_1st, sigma_long_1st


    def get_mean_sigma(self, t):
        """
        由初始时刻的平均根数得到t时间后的瞬时平均根数
        :param t: 时间
        :return: 瞬时平均根数
        """

        sigma_short_1st, sigma_long_1st = self.get_1st_short_and_long()
        mean_sigma_0 = self.sigma - sigma_short_1st - sigma_long_1st
        # 保存一下现在的sigma
        sigma_0 = self.sigma
        # 更新一下，用于下面重新计算a的1阶短周期项
        self.__init__(mean_sigma_0)
        # 重新计算a的短周期项，使其精确到2阶
        sigma_short_2nd, sigma_long_2nd = self.get_1st_short_and_long()
        sigma_short_1st[0] = sigma_short_2nd[0]
        # 获取初始时刻的平均根数
        mean_sigma_0_a_short_is_2nd = sigma_0 - sigma_short_1st - sigma_long_1st

        a, e, i, omega, w, M = mean_sigma_0_a_short_is_2nd

        E = utils.newton_generation(M, e, M, 1e-12)
        r = a * (1 - e * np.cos(E))
        p = a * (1 - e ** 2)
        f = 2 * np.arctan(np.sqrt(1 + e) / (1 - e)) * np.tan(E / 2)
        n = a ** (-3/2)

        a_1, e_1, i_1 = [0, 0, 0]
        omega_1 = -A_2 / (p ** 2) * n * np.cos(i)
        w_1 = A_2 / (p ** 2) * n * (2 - 5/2 * (np.sin(i) ** 2))
        M_1 = A_2 / (p ** 2) * n * (1 - 3/2 * (np.sin(i) ** 2)) * np.sqrt(1 - e ** 2)
        sigma_1 = np.array([a_1, e_1, i_1, omega_1, w_1, M_1])

        a_2, e_2, i_2 = [0, 0, 0]
        omega_2 = - (A_2 / (p * p)) ** 2 * n * np.cos(i) * ((3/2 + 1/6 * e * e + np.sqrt(1 - e * e)) - (5/3 - 5/24 * e * e + 3/2 * np.sqrt(1 - e * e)) * (np.sin(i) ** 2))
        w_2 = (A_2 / (p * p)) ** 2 * n * ((4 + 7/12 * e * e + 2 * np.sqrt(1 - e ** 2)) - (103/12 + 3/8 * e * e + 11/2 * np.sqrt(1 - e * e)) * (np.sin(i) ** 2) + (215/48 - 15/32 * e * e + 15/4 * np.sqrt(1 - e * e)) * (np.sin(i) ** 4))
        M_2 = (A_2 / (p * p)) ** 2 * n * np.sqrt(1 - e * e) * (1/2 * (1 - 3/2 * (np.sin(i) ** 2)) ** 2 * np.sqrt(1 - e * e) + 5/2 + 10/3 * e * e - (19/3 + 26/3 * e * e) * (np.sin(i) ** 2) + (233/48 + 103/12 * e * e) * (np.sin(i) ** 4) + (e ** 4) / (1 - e * e) * (35/12 - 35/4 * (np.sin(i) ** 2) + 315/32 * (np.sin(i) ** 4)))
        sigma_2 = np.array([a_2, e_2, i_2, omega_2, w_2, M_2])

        mean_sigma = mean_sigma_0_a_short_is_2nd + (np.array([0, 0, 0, 0, 0, n]) + sigma_1 + sigma_2) * t
        return mean_sigma

    def get_final_sigma(self, t):
        """
        :param t: 时间
        :return: 瞬时根数
        """

        # 计算瞬时平均根数
        mean_sigma = self.get_mean_sigma(t)
        # 更新根数，用于计算长短周期项
        self.__init__(mean_sigma)
        # 计算瞬时根数
        mean_sigma_short_1st, mean_sigma_long_1st = self.get_1st_short_and_long()
        final_sigma = mean_sigma + mean_sigma_short_1st + mean_sigma_long_1st
        return final_sigma


class ValueIntegration:
    def __init__(self, sigma=np.zeros(6), h=10, time=0):
        self.sigma = sigma
        self.a = sigma[0]
        self.e = sigma[1]
        self.i = sigma[2]
        self.omega = sigma[3]
        self.w = sigma[4]
        self.M = sigma[5]
        self.E = utils.newton_generation(self.M, self.e, self.M, 1e-12)
        self.r = self.a * (1 - self.e * np.cos(self.E))
        self.p = self.a * (1 - self.e ** 2)
        self.n = self.a ** (-3/2)
        # self.f = 2 * np.arctan(np.sqrt(1 + self.e) / (1 - self.e)) * np.tan(self.E / 2)
        self.f = np.arctan2(np.sqrt(1 - self.e ** 2) * np.sin(self.E), np.cos(self.E) - self.e)
        self.h = h
        self.t = time

    def right_function(self, sigma, t):

        """
        给定sigma，按照公式计算对应的右函数
        :param sigma:轨道根数
        :param t:时间,没用到
        :return:返回对应的右函数
        """
        # 用传入的sigma进行更新
        self.__init__(sigma)

        # 按照书本（3.60）式计算S, T, W
        S = -A_2 / (self.r ** 4) * (1 - 1.5 * (np.sin(self.i) ** 2) + 1.5 * (np.sin(self.i) ** 2) * np.cos(2 * (self.f + self.w)))
        T = -A_2 / (self.r ** 4) * (np.sin(self.i) ** 2) * np.sin(2 * (self.f + self.w))
        W = A_2 / (self.r ** 4) * np.sin(2 * self.i) * (-np.sin(self.f + self.w))

        # 按照书本（3.73）式计算根数对t的导数
        diff_a = 2 / (self.n * np.sqrt(1 - self.e ** 2)) * (S * self.e * np.sin(self.f) + T * (1 + self.e * np.cos(self.f)))
        diff_e = np.sqrt(1 - self.e ** 2) / (self.n * self.a) * (S * np.sin(self.f) + T * (np.cos(self.f) + np.cos(self.E)))
        diff_i = self.r * np.cos(self.f + self.w) / (self.n * self.a * self.a * np.sqrt(1 - self.e ** 2)) * W
        diff_omega = self.r * np.sin(self.f + self.w) / (self.n * self.a * self.a * np.sqrt(1 - self.e ** 2) * np.sin(self.i)) * W
        diff_w = np.sqrt(1 - self.e ** 2) / (self.n * self.a * self.e) * (-S * np.cos(self.f) + T * (1 + self.r / self.p) * np.sin(self.f)) - np.cos(self.i) * diff_omega
        diff_M = self.n - (1 - self.e ** 2) / (self.n * self.a * self.e) * (-S * (np.cos(self.f) - 2 * self.e * self.r / self.p) + T * (1 + self.r / self.p) * np.sin(self.f))

        return np.array([diff_a, diff_e, diff_i, diff_omega, diff_w, diff_M])

    def value_integration(self):
        """
        对微分方程进行数值积分
        :return:积分后的值
        """
        sigma = self.sigma
        # 对步长进行归一化
        stride = self.h / time_unit
        for i in range(int(self.t / self.h)):
            # 固定w,M在[0,2pi)内
            sigma[4] = sigma[4] % (2 * np.pi)
            sigma[5] = sigma[5] % (2 * np.pi)
            # 使用rk4进行积分
            sigma = utils.rk4(sigma, 0, stride, self.right_function)

        return sigma


# 初始轨道根数
sigma_0 = np.array([6800000 / a_e, 0.01, np.pi / 4, 0, 0, 0])
# 计算时间，进行归一化
time = 24 * 3600
time_standard = time / time_unit

# 平均根数法计算轨道根数
me = MeanElements(sigma_0)
me_final_sigma = me.get_final_sigma(time_standard)
# 逆归一化
me_final_sigma[0] = me_final_sigma[0] * a_e
me_final_sigma[2:] = me_final_sigma[2:] * 180 / np.pi % 360

# 数值积分法计算轨道根数
vi = ValueIntegration(sigma_0, 10, time)
vi_final_sigma = vi.value_integration()
# 逆归一化
vi_final_sigma[0] = vi_final_sigma[0] * a_e
vi_final_sigma[2:] = vi_final_sigma[2:] * 180 / np.pi % 360

# 控制输出格式，方便对比
np.set_printoptions(formatter={'float': '{:.6f}'.format})
print("平均根数法")
print(me_final_sigma)
print("数值积分法")
print(vi_final_sigma)

# me_final_sigma_list = open('me_final_sigma_list.txt', 'w')
# vi_final_sigma_list = open('vi_final_sigma_list.txt', 'w')
# difference_sigma_list = open('difference_sigma_list.txt', 'w')
#
# t_0 = '2013-09-01T12:00:00'
#
# for i in range(24 * 60):
#     time = i * 60 + 60
#     time_standard = time / time_unit
#     t = Time(t_0, scale='utc') + time * u.s
#     # 平均根数法计算轨道根数
#     me = MeanElements(sigma_0)
#     me_final_sigma = me.get_final_sigma(time_standard)
#
#     # 数值积分法计算轨道根数
#     vi = ValueIntegration(sigma_0, 10, time)
#     vi_final_sigma = vi.value_integration()
#
#     difference_sigma = vi_final_sigma - me_final_sigma
#
#     me_final_sigma_list.write(str(t) + str(me_final_sigma) + '\n')
#     vi_final_sigma_list.write(str(t) + str(vi_final_sigma) + '\n')
#     difference_sigma_list.write(str(t) + str(difference_sigma) + '\n')
#
# print("ok")
# me_final_sigma_list.close()
# vi_final_sigma_list.close()
# difference_sigma_list.close()